'''
img1 img2
depth(img1) + emb(img2) = pesudo img w/ img1's shape 
mask(pesudo) + fill(pesudo) + emb(img2) = final img
'''

import torch,os,pdb
from diffusers import FluxFillPipeline,FluxPriorReduxPipeline,FluxControlPipeline
from diffusers.utils import load_image
from itertools import product
from image_gen_aux import DepthPreprocessor

from util_flux import pad_image
from util_flux import horizontal_concat_images

from util_sam import get_mask_by_sam
from util_mask import get_erosed_mask_by_radtio,add_random_holes


FLUX_FILL='/home/shengjie/ckp/FLUX.1-Fill-dev'
FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_DEPTH='/home/shengjie/ckp/FLUX.1-Depth-dev'
FLUX_DEPTH_LORA='/home/shengjie/ckp/FLUX.1-Depth-dev-lora'
FLUX='/data/models/FLUX___1-dev'

DEPTH_PREDCITION='/home/shengjie/ckp/depth-anything-large-hf'


target_shape = (1024,1024)

examples_dir = '/data/shengjie/style_zhenzhi/'
## will save
save_dir = '/data/shengjie/synthesis_zhenzhi2/'
depth_dir = '/data/shengjie/style_zhenzhi_depth/'
embding_redux_dir = '/data/shengjie/style_zhenzhi_emb/'
pesudo_dir = '/data/shengjie/style_zhenzhi_pesudo/'
os.makedirs(save_dir,exist_ok=True)
os.makedirs(depth_dir,exist_ok=True)
os.makedirs(embding_redux_dir,exist_ok=True)
os.makedirs(pesudo_dir,exist_ok=True)

imagefiles = os.listdir(examples_dir)

from tqdm import tqdm
for t1,t2 in tqdm(product(imagefiles,imagefiles)):
    if t1 == t2:continue

    img1_path = os.path.join(examples_dir,t1)
    img2_path = os.path.join(examples_dir,t2)
    # PIL
    img1 = load_image(img1_path) # PIL
    img1,_,_,_,_ = pad_image(img1)
    img1 = img1.resize(target_shape)
    img2 = load_image(img2_path) # PIL
    from util_mask import center_crop
    # img2 = center_crop(img2)
    img2,_,_,_,_ = pad_image(img2)
    img2 = img2.resize(target_shape)

    processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)
    ## depth(img1)
    img1_depth = processor(img1)[0].convert("RGB")
    del processor

    ## emb(img2)
    pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                                FLUX_REDUX, 
                                                torch_dtype=torch.bfloat16).to("cuda")
    main_condition_prompt = pipe_prior_redux(img2) # attr 'prompt_embeds' torch.Size([1, 1241, 4096]) 
    del pipe_prior_redux

    pipe = FluxControlPipeline.from_pretrained(FLUX, torch_dtype=torch.bfloat16).to("cuda")
    pipe.load_lora_weights(FLUX_DEPTH_LORA, adapter_name="depth")
    pipe.set_adapters("depth", 0.85)

    ## pesudo
    pesudo = pipe(
            control_image=img1_depth,
            height=target_shape[1],
            width=target_shape[0],
            num_inference_steps=8,
            guidance_scale=10.0,
            # generator=torch.Generator().manual_seed(42),
            **main_condition_prompt
        ).images[0]
    del pipe

    torch.cuda.empty_cache()
    # pdb.set_trace()

    ## mask(pesudo)
    mask_pesudo_ori,mask_pesudo_pil_ori = get_mask_by_sam(image_pil=pesudo)
    # mask,mask_pil = get_erosed_mask_by_radtio(mask,0.5)

    pipe = FluxFillPipeline.from_pretrained(FLUX_FILL, torch_dtype=torch.bfloat16).to("cuda")


    for hole_size,hole_nums in zip([20,30,40,50],[180,80,50,30]):
    # for hole_size,hole_nums in product([30,40,50],[30,40,50]):

    # hole_size,hole_nums = 100,10

        ## add holes to mask(pesudo)
        mask_hole , mask_hole_pil = add_random_holes(mask_pesudo_pil_ori,
                                                    min_holes=hole_nums,
                                                    max_holes=hole_nums,
                                                    hole_size_range=(hole_size,hole_size))
        # pdb.set_trace()
        mask_pil = mask_hole_pil
        mask,_,_,_,_ = pad_image(mask_pil,fill=(255,))
        mask = mask.resize(target_shape)
        pesudo_mask = mask

        # pdb.set_trace()

        ## final img  
        res_image = pipe(
            # prompt="a white paper cup",
            image=pesudo,
            mask_image=pesudo_mask,
            height=target_shape[1],
            width=target_shape[0],
            guidance_scale=30,
            num_inference_steps=8,
            max_sequence_length=512,
            # generator=torch.Generator("cpu").manual_seed(0),
            **main_condition_prompt,
        ).images[0]

        concat_res = horizontal_concat_images([img1,img1_depth,img2,pesudo,pesudo_mask,res_image])

        ## save 
        save_path = os.path.join(save_dir,
                                     f"{os.path.splitext(t1)[0]}_{os.path.splitext(t2)[0]}"+\
                                    f'_hsize{hole_size}_hnums{hole_nums}{os.path.splitext(t1)[1]}'
                                    )
        save_name = f'./tmp_hsize{hole_size}_hnums{hole_nums}{os.path.splitext(t1)[1]}'
        concat_res.save(save_name)
        pdb.set_trace()
        concat_res.save(save_path)
    del pipe
    torch.cuda.empty_cache()
